Skip to content

Sequence parallelism with torch rms norm #27062

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Closed
wants to merge 54 commits into from
Closed

Sequence parallelism with torch rms norm #27062

wants to merge 54 commits into from

Conversation

@ProExpertProg
Copy link
Collaborator

@ProExpertProg ProExpertProg commented Oct 17, 2025 •
edited by github-actions bot
Loading

This is on top of 2.9 + #24604 as I didn't want any surprises so I just developed with 2.9. We can check 2.8 with CI :P.

TODOs:

  • I think it would be nice to merge Middle and Last patterns - see allreduce fusion in [torch.compile] Enable attention and allreduce fusion without custom ops enabled #24604 (we can wrap the pattern and replacement to just return the first output)
  • I think the sequence parallelism unit test is kinda complicated with fusion and fix_finctionalization happening. rmsnorm+quant fusion stopped working for me, if you can get it working that's good, if not no worries.

huydhn added 30 commits September 16, 2025 13:55
@huydhn
Update PyTorch to 2.9.0
e738cba 
Signed-off-by: Huy Do <huydhn@gmail.com>
@huydhn
Add a comment
497cffd 
Signed-off-by: Huy Do <huydhn@gmail.com>
@huydhn
Not setting --extra-index-url in test.in
89ba43a 
Signed-off-by: Huy Do <huydhn@gmail.com>
@huydhn
Use https://download.pytorch.org/whl/test
157aae3 
Signed-off-by: Huy Do <huydhn@gmail.com>
@huydhn
Merge branch 'main' into pytorch-2.9.0
ea1ef7a 
Signed-off-by: Huy Do <huydhn@gmail.com>
@huydhn
Put torchao back to the same state
0a81fb2 
Signed-off-by: Huy Do <huydhn@gmail.com>
@huydhn
Install the latest torchao nightly for quantization test
0325966 
Signed-off-by: Huy Do <huydhn@gmail.com>
@huydhn
Debug distributed failures
39b9cbf 
Signed-off-by: Huy Do <huydhn@gmail.com>
@huydhn
Wrong torchao package
0272040 
Signed-off-by: Huy Do <huydhn@gmail.com>
@huydhn
Attempt the fix in NVIDIA/nccl#1838
c2e0eaf 
Signed-off-by: Huy Do <huydhn@gmail.com>
@huydhn
Merge branch 'main' into pytorch-2.9.0
c16db74 
Signed-off-by: Huy Do <huydhn@gmail.com>
@huydhn
Set inductor_graph_partition to True by default
d3436a8 
Signed-off-by: Huy Do <huydhn@gmail.com>
@huydhn
Rerun with RC3
0e581a3 
Signed-off-by: Huy Do <huydhn@gmail.com>
@huydhn
Rerun with RC4
84c6cc3 
Signed-off-by: Huy Do <huydhn@gmail.com>
@huydhn
Merge branch 'main' into pytorch-2.9.0
3637adb 
Signed-off-by: Huy Do <huydhn@gmail.com>
@huydhn
Build CPU docker image
23c6427 
Signed-off-by: Huy Do <huydhn@gmail.com>
@huydhn
Leave CPU for later
ba8a85f 
Signed-off-by: Huy Do <huydhn@gmail.com>
@huydhn
CPU build should work now
ec7b5c4 
Signed-off-by: Huy Do <huydhn@gmail.com>
@huydhn
Rebuild flashinfer-python for 2.9.0
869d13e 
Signed-off-by: Huy Do <huydhn@gmail.com>
@huydhn
Merge branch 'main' into pytorch-2.9.0
a670c2e 
Signed-off-by: Huy Do <huydhn@gmail.com>
@huydhn
Fix precommit
145e225 
Signed-off-by: Huy Do <huydhn@gmail.com>
@huydhn
Merge branch 'main' into pytorch-2.9.0
9cd7683 
Signed-off-by: Huy Do <huydhn@gmail.com>
@huydhn
Merge branch 'main' into pytorch-2.9.0
47ae5d8 
Signed-off-by: Huy Do <huydhn@gmail.com>
@huydhn
Merge branch 'main' into pytorch-2.9.0
e7064b4 
Signed-off-by: Huy Do <huydhn@gmail.com>
@huydhn
Merge branch 'main' into pytorch-2.9.0
106bd40 
Signed-off-by: Huy Do <huydhn@gmail.com>
@huydhn
Merge branch 'main' into pytorch-2.9.0
76e438d 
Signed-off-by: Huy Do <huydhn@gmail.com>
@huydhn
Merge branch 'main' into pytorch-2.9.0
ebaa419 
Signed-off-by: Huy Do <huydhn@gmail.com>
@huydhn
Merge branch 'main' into pytorch-2.9.0
b4ed78c 
Signed-off-by: Huy Do <huydhn@gmail.com>
@huydhn
Skip some test unless it's B200
5d50c59 
Signed-off-by: Huy Do <huydhn@gmail.com>
@huydhn
Merge branch 'main' into pytorch-2.9.0
210aa68 
Signed-off-by: Huy Do <huydhn@gmail.com>
angelayi and others added 11 commits October 17, 2025 14:35
@angelayi @ProExpertProg
[cherry-pick] [compile] Enable sequence parallelism matching w/o cust…
8552681 
…om ops enabled

Signed-off-by: angelayi <yiangela7@gmail.com>

Signed-off-by: ProExpertProg <lgovedic@redhat.com>
@ProExpertProg
refactor e2e test to use Matches object
4ac901e 
Signed-off-by: ProExpertProg <lgovedic@redhat.com>
@ProExpertProg
Add dynamic shape to seqpar tests
00d63e9 
Signed-off-by: ProExpertProg <lgovedic@redhat.com>
@ProExpertProg
Add E2E test for SeqPar+asyncTP (TODO rms_norm still needed)
101f75d 
Signed-off-by: ProExpertProg <lgovedic@redhat.com>
@ProExpertProg
Skip E2E test for inductor-partition + fp4
9dadf9a 
Signed-off-by: ProExpertProg <lgovedic@redhat.com>
@ProExpertProg
Add rmsnorm to seqpar unit tests
3f69d07 
Signed-off-by: ProExpertProg <lgovedic@redhat.com>
@ProExpertProg
remove rms_norm from vllm config for seq par
05993dd 
Signed-off-by: ProExpertProg <lgovedic@redhat.com>
@ProExpertProg
Fix sequence parallelism for torch rmsnorm
b63e40d 
Signed-off-by: ProExpertProg <lgovedic@redhat.com>
@ProExpertProg
TEMP test sequence parallelism
e2d120c 
Signed-off-by: ProExpertProg <lgovedic@redhat.com>
@ProExpertProg
Add comment
8e15e19 
Signed-off-by: ProExpertProg <lgovedic@redhat.com>
@ProExpertProg
E2E test with rms_norm
88f9e30 
Signed-off-by: ProExpertProg <lgovedic@redhat.com>
@ProExpertProg ProExpertProg force-pushed the luka/2.9-inductor-partition-seq-par-no-rms-norm branch from e6b49b4 to 88f9e30 Compare October 17, 2025 14:37
@ProExpertProg ProExpertProg marked this pull request as ready for review October 17, 2025 14:43
@ProExpertProg ProExpertProg mentioned this pull request Oct 20, 2025
Open
@mergify
Copy link

mergify bot commented Oct 23, 2025

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @ProExpertProg.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

@mergify mergify bot added the needs-rebase label Oct 23, 2025
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@bigPYJ1151 bigPYJ1151 Awaiting requested review from bigPYJ1151 bigPYJ1151 is a code owner

@hmellor hmellor Awaiting requested review from hmellor hmellor is a code owner

@jeejeelee jeejeelee Awaiting requested review from jeejeelee jeejeelee is a code owner

@simon-mo simon-mo Awaiting requested review from simon-mo simon-mo is a code owner

@WoosukKwon WoosukKwon Awaiting requested review from WoosukKwon WoosukKwon is a code owner

@youkaichao youkaichao Awaiting requested review from youkaichao youkaichao is a code owner

@robertgshaw2-redhat robertgshaw2-redhat Awaiting requested review from robertgshaw2-redhat robertgshaw2-redhat is a code owner

@mgoin mgoin Awaiting requested review from mgoin mgoin is a code owner

@tlrmchlsmth tlrmchlsmth Awaiting requested review from tlrmchlsmth tlrmchlsmth is a code owner

@houseroad houseroad Awaiting requested review from houseroad houseroad is a code owner

@yewentao256 yewentao256 Awaiting requested review from yewentao256 yewentao256 is a code owner

@LucasWilkinson LucasWilkinson Awaiting requested review from LucasWilkinson LucasWilkinson is a code owner

@zou3519 zou3519 Awaiting requested review from zou3519 zou3519 is a code owner

Assignees

No one assigned

Labels

ci/build documentation Improvements or additions to documentation needs-rebase ready ONLY add when PR is ready to merge/full CI is needed rocm Related to AMD ROCm

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

3 participants

@ProExpertProg @huydhn @angelayi